"""
扣子AI智能问答API接口
使用扣子平台的流式接口进行智能问答
"""
import os
import json
import uuid
import logging
from typing import Optional
from fastapi import APIRouter, HTTPException, Request, status
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from utils import JWTUtils

# 导入扣子SDK
try:
    from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageContentType, ChatEventType, COZE_CN_BASE_URL
    COZE_AVAILABLE = True
except ImportError:
    COZE_AVAILABLE = False
    logging.warning("cozepy未安装，扣子AI功能将不可用。请运行: pip install cozepy")

router = APIRouter(prefix="/coze", tags=["扣子AI"])

# 扣子API配置
COZE_API_TOKEN = os.getenv("COZE_API_TOKEN", "pat_ujR7mHN5GAcwZqCqowV8Lf6hCFhH3i6r00XDnvshUAgWh4pJnrtB5CdARRfMPf8X")
COZE_BOT_ID = os.getenv("COZE_BOT_ID", "7569434809961971753")  # 默认bot_id，可以从环境变量配置

# 初始化扣子客户端
if COZE_AVAILABLE:
    try:
        coze = Coze(
            auth=TokenAuth(token=COZE_API_TOKEN),
            base_url=COZE_CN_BASE_URL
        )
    except Exception as e:
        logging.error(f"初始化扣子客户端失败: {e}")
        coze = None
else:
    coze = None

logger = logging.getLogger(__name__)


class CozeChatRequest(BaseModel):
    """扣子AI问答请求"""
    message: str
    user_id: Optional[str] = None
    conversation_id: Optional[str] = None  # 对话ID，用于多轮对话


class CozeChatResponse(BaseModel):
    """扣子AI问答响应"""
    message: str
    user_id: str
    conversation_id: Optional[str] = None
    token_usage: Optional[dict] = None


# 流式聊天接口
@router.post("/chat/stream", summary="扣子AI流式问答")
async def coze_chat_stream(request_data: CozeChatRequest, request: Request):
    """
    扣子AI流式问答接口
    返回流式响应，实时返回AI回复内容
    """
    if not COZE_AVAILABLE or not coze:
        raise HTTPException(
            status_code=503,
            detail="扣子AI服务不可用，请检查cozepy是否已安装"
        )
    
    if not request_data.message or not request_data.message.strip():
        raise HTTPException(status_code=400, detail="请输入问题")
    
    # 获取或生成用户ID
    user_id = request_data.user_id
    if not user_id:
        user_id = str(uuid.uuid4())
    
    try:
        async def generate():
            """生成器函数，用于流式返回"""
            full_message = ""
            conversation_id = request_data.conversation_id
            token_usage = None
            
            try:
                # 调用扣子流式接口
                for event in coze.chat.stream(
                    bot_id=COZE_BOT_ID,
                    user_id=user_id,
                    conversation_id=conversation_id,
                    additional_messages=[
                        Message.build_user_question_text(request_data.message),
                    ],
                ):
                    if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
                        # 增量消息内容
                        content = event.message.content
                        if content:
                            full_message += content
                            # 发送增量内容
                            yield f"data: {json.dumps({'type': 'delta', 'content': content}, ensure_ascii=False)}\n\n"
                    
                    elif event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
                        # 对话完成
                        conversation_id = event.chat.conversation_id if hasattr(event.chat, 'conversation_id') else None
                        if hasattr(event.chat, 'usage') and event.chat.usage:
                            token_usage = {
                                'token_count': event.chat.usage.token_count if hasattr(event.chat.usage, 'token_count') else None
                            }
                        # 发送完成消息
                        yield f"data: {json.dumps({'type': 'completed', 'message': full_message, 'conversation_id': conversation_id, 'token_usage': token_usage}, ensure_ascii=False)}\n\n"
                    
                    elif event.event == ChatEventType.CONVERSATION_CHAT_FAILED:
                        # 错误事件
                        error_msg = str(event.error) if hasattr(event, 'error') and event.error else "未知错误"
                        yield f"data: {json.dumps({'type': 'error', 'error': error_msg}, ensure_ascii=False)}\n\n"
                        break
                
            except Exception as e:
                logger.error(f"扣子AI流式调用失败: {e}", exc_info=True)
                yield f"data: {json.dumps({'type': 'error', 'error': str(e)}, ensure_ascii=False)}\n\n"
        
        return StreamingResponse(
            generate(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no"  # 禁用nginx缓冲
            }
        )
    
    except Exception as e:
        logger.error(f"扣子AI请求处理失败: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")


# 非流式聊天接口（兼容性接口）
@router.post("/chat", response_model=CozeChatResponse, summary="扣子AI问答")
async def coze_chat(request_data: CozeChatRequest, request: Request):
    """
    扣子AI问答接口（非流式）
    等待完整回复后返回
    """
    if not COZE_AVAILABLE or not coze:
        raise HTTPException(
            status_code=503,
            detail="扣子AI服务不可用，请检查cozepy是否已安装"
        )
    
    if not request_data.message or not request_data.message.strip():
        raise HTTPException(status_code=400, detail="请输入问题")
    
    # 获取或生成用户ID
    user_id = request_data.user_id
    if not user_id:
        user_id = str(uuid.uuid4())
    
    try:
        full_message = ""
        conversation_id = request_data.conversation_id
        token_usage = None
        
        # 调用扣子流式接口，但收集完整消息
        for event in coze.chat.stream(
            bot_id=COZE_BOT_ID,
            user_id=user_id,
            conversation_id=conversation_id,
            additional_messages=[
                Message.build_user_question_text(request_data.message),
            ],
        ):
            if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
                content = event.message.content
                if content:
                    full_message += content
            
            elif event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
                conversation_id = event.chat.conversation_id if hasattr(event.chat, 'conversation_id') else None
                if hasattr(event.chat, 'usage') and event.chat.usage:
                    token_usage = {
                        'token_count': event.chat.usage.token_count if hasattr(event.chat.usage, 'token_count') else None
                    }
                break
            
            elif event.event == ChatEventType.CONVERSATION_CHAT_FAILED:
                error_msg = str(event.error) if hasattr(event, 'error') and event.error else "未知错误"
                raise HTTPException(status_code=500, detail=f"扣子AI错误: {error_msg}")
        
        if not full_message:
            raise HTTPException(status_code=500, detail="未收到有效回复")
        
        return CozeChatResponse(
            message=full_message,
            user_id=user_id,
            conversation_id=conversation_id,
            token_usage=token_usage
        )
    
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"扣子AI请求处理失败: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")

